package similarity

import (
	"embed"
	"errors"
	"fmt"
	"gitee.com/jn-qq/big-model/common"
	"gitee.com/jn-qq/big-model/openai"
	"github.com/go-ego/gse"
	"log"
	"math"
	"os"
	"strconv"
	"strings"
)

//go:embed dict
var dict embed.FS

type Similar struct {
	seg gse.Segmenter
}

// SimilarMethod 相似度计算方法
type SimilarMethod int

const (
	// Cosine 余弦相似度
	Cosine SimilarMethod = iota
	// CosineTFIDF TF-IDF加权余弦相似度
	CosineTFIDF
	// Jaccard Jaccard相似度
	Jaccard
	// LevenshteinDistance 编辑距离
	LevenshteinDistance
	// SimHashDistance SimHash相似度
	SimHashDistance
	// ModelSimilarity 模型相似度
	ModelSimilarity
)

// NewSimilar 创建相似度计算对象
func NewSimilar(files ...string) (*Similar, error) {
	if len(files) == 0 {
		files = append(files, "zh_s")
	}
	// 创建分词器
	seg, err := gse.New(files...)
	if err != nil {
		return nil, err
	}

	return &Similar{seg: seg}, nil
}

// LoadStopWords 加载停用词
func (s *Similar) LoadStopWords(files ...string) error {
	if len(files) == 0 {
		log.Println("无指定停用词表，跳过")
		return nil
	} else {
		log.Println("加载停用词表:", files)
	}

	// 创建停用词存储对象
	s.seg.StopWordMap = make(map[string]bool)

	// 加载默认停用词,如果有
	if files[0] == "DEFAULT" {
		files = files[1:]
		log.Println("加载默认停用词表")
		stopFiles, err := dict.ReadDir("dict/stop_words")
		if err != nil {
			return err
		}
		for _, file := range stopFiles {
			fileBytes, err := dict.ReadFile("dict/stop_words/" + file.Name())
			if err != nil {
				return err
			}
			for _, word := range strings.Split(string(fileBytes), "\n") {
				s.seg.AddStop(word)
			}
		}
	}

	// 加载自定义停用词
	for _, file := range files {
		fileBytes, err := os.ReadFile(file)
		if err != nil {
			return err
		}
		for _, word := range strings.Split(string(fileBytes), "\n") {
			s.seg.AddStop(word)
		}
	}
	return nil
}

// Tokenize 文本清洗与分词
func (s *Similar) Tokenize(text string) []string {
	// 1. 基本清洗
	text = strings.TrimSpace(text)
	//text = regexp.MustCompile(`[^\p{Han}a-zA-Z0-9]`).ReplaceAllString(text, " ")
	//text = regexp.MustCompile(`\s+`).ReplaceAllString(text, " ")
	text = strings.ToLower(text)

	// 2. 中文分词
	words := s.seg.Cut(text, true) // 精确模式

	// 3. 过滤停用词和单字
	var tokens []string
	for _, word := range words {
		word = strings.TrimSpace(word)
		if s.seg.IsStop(word) {
			continue
		}
		tokens = append(tokens, word)
	}
	return tokens
}

// cosineSimilarity 计算两个文本的余弦相似度
func (s *Similar) cosineSimilarity(text1, text2 string, useTFIDF bool) float64 {
	tokens1 := s.Tokenize(text1)
	tokens2 := s.Tokenize(text2)
	//fmt.Println(tokens1, tokens2)

	// 构建词频向量
	vec1 := make(map[string]float64)
	vec2 := make(map[string]float64)

	if useTFIDF {
		// 使用TF-IDF加权
		// 构造文档集合
		documents := documentSet{
			Document: [][]string{tokens1, tokens2},
		}
		// 计算IDF
		documents.calculateIDF()
		// 计算每个文档的TF
		documents.calculateTF()
		// 计算TF-IDF
		tfidfList := documents.calculateTFIDF()
		vec1 = tfidfList[0]
		vec2 = tfidfList[1]
	} else {
		// 使用原始词频
		for _, token := range tokens1 {
			vec1[token]++
		}

		for _, token := range tokens2 {
			vec2[token]++
		}
	}
	//fmt.Println(vec1, vec2)

	// 计算点积
	var dotProduct float64
	for token, freq1 := range vec1 {
		if freq2, exists := vec2[token]; exists {
			dotProduct += freq1 * freq2
		}
	}

	// 计算向量的模
	var magnitude1, magnitude2 float64
	for _, freq := range vec1 {
		magnitude1 += freq * freq
	}
	for _, freq := range vec2 {
		magnitude2 += freq * freq
	}

	magnitude1 = math.Sqrt(magnitude1)
	magnitude2 = math.Sqrt(magnitude2)

	// 避免除零
	if magnitude1 == 0 || magnitude2 == 0 {
		return 0
	}

	return dotProduct / (magnitude1 * magnitude2)
}

// jaccardSimilarity 计算两个文本的Jaccard相似度
func (s *Similar) jaccardSimilarity(text1, text2 string) float64 {
	tokens1 := s.Tokenize(text1)
	tokens2 := s.Tokenize(text2)

	// 转换为集合
	set1 := make(map[string]struct{})
	set2 := make(map[string]struct{})

	for _, token := range tokens1 {
		set1[token] = struct{}{}
	}

	for _, token := range tokens2 {
		set2[token] = struct{}{}
	}

	// 计算交集
	intersection := make(map[string]struct{})
	for token := range set1 {
		if _, exists := set2[token]; exists {
			intersection[token] = struct{}{}
		}
	}

	// 计算并集
	union := make(map[string]struct{})
	for token := range set1 {
		union[token] = struct{}{}
	}
	for token := range set2 {
		union[token] = struct{}{}
	}

	if len(union) == 0 {
		return 0
	}

	return float64(len(intersection)) / float64(len(union))
}

// LevenshteinDistance 编辑距离
func (s *Similar) levenshteinDistance(s1, s2 string) float64 {
	runes1 := []rune(s1)
	runes2 := []rune(s2)
	len1 := len(runes1)
	len2 := len(runes2)

	// 只需要两行来节省空间
	prev := make([]int, len2+1)
	curr := make([]int, len2+1)

	// 初始化第一行
	for j := 0; j <= len2; j++ {
		prev[j] = j
	}

	for i := 1; i <= len1; i++ {
		curr[0] = i
		for j := 1; j <= len2; j++ {
			cost := 1
			if runes1[i-1] == runes2[j-1] {
				cost = 0
			}

			curr[j] = min(
				prev[j]+1,      // 删除
				curr[j-1]+1,    // 插入
				prev[j-1]+cost, // 替换或匹配
			)
		}
		// 交换prev和curr
		prev, curr = curr, prev
	}

	maxLen := max(len1, len2)
	if maxLen == 0 {
		return 1.0
	}
	return 1.0 - float64(prev[len2])/float64(maxLen)
}

var ModelConfig common.Model

func (s *Similar) modelSimilarity(s1, s2 string) (float64, error) {
	if ModelConfig.Url == "" {
		return math.NaN(), errors.New("ModelConfig is nil")
	}
	// 构建提示词
	prompt := "<User>你是一个专业的文本相似度评估助手。请分析以下两段文本的语义相似度，返回一个0到1之间的数字，不要有任何其他文字或解释。" +
		"<ret>## 规则：<ret>- 0表示两段文本语义完全不同<ret>- 1表示两段文本语义完全相同。<ret>- 请只输出得分数字（保留2位小数），不要添加任何其他文字或解释。<ret>" +
		"<ret>## 待分析文本<ret>- 文本1：%s<ret>- 文本2：%s<ret>" +
		"<ret>## 请返回相似度(0-1)：<end><Bot>"
	prompt = fmt.Sprintf(prompt, s1, s2)
	chat := openai.NewOpenAiClient(ModelConfig).
		SetMaxTokens(common.DefaultTokens).
		SetTemperature(common.DefaultTemperature)
	response, err := chat.Chat(prompt)
	if err != nil {
		return math.NaN(), err
	}
	answer := response.GetAnswer()
	// 转为float64
	similarity, err := strconv.ParseFloat(answer, 64)
	if err != nil {
		return math.NaN(), fmt.Errorf("strconv.ParseFloat err: %v", err)
	}
	return similarity, nil
}

// TextSimilarity 计算两个文本的综合相似度
func (s *Similar) TextSimilarity(text1, text2 string, similarMethod ...SimilarMethod) (float64, error) {
	similarMethod = append(similarMethod, Cosine)
	switch similarMethod[0] {
	case Cosine:
		return s.cosineSimilarity(text1, text2, false), nil

	case CosineTFIDF:
		return s.cosineSimilarity(text1, text2, true), nil

	case Jaccard:
		return s.jaccardSimilarity(text1, text2), nil

	case LevenshteinDistance:
		return s.levenshteinDistance(text1, text2), nil
	//
	//case NGram:
	//	return ngram.Similarity(text1, text2)
	//
	case SimHashDistance:
		_simHash := newSimHash()
		return _simHash.SimHashSimilarity(s.Tokenize(text1), s.Tokenize(text2)), nil

	case ModelSimilarity:
		return s.modelSimilarity(text1, text2)
	default:
		return math.NaN(), errors.New("invalid similarity method")
	}
}
